from langchain.agents import create_react_agent, AgentExecutor
from langchain_chroma import Chroma
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool

from home.models import HistoryModel


def search_db(input):
    embeddings = DashScopeEmbeddings()
    db = Chroma(persist_directory="./home/database/chroma_db", embedding_function=embeddings)
    res = db.similarity_search(input, k=3)
    data = ""
    for i in res:
        data += i.page_content + "\n"
    return data


llm = Tongyi()


def tool1(input):
    data = HistoryModel.objects.all().order_by('-created_at')[:3]
    datas = [i.history for i in data]
    print(datas)
    return datas


def tool2(input):
    return "曾经咨询过类似的产品"


tools = [
    Tool(func=tool1, name="tool1", description="对新产品感兴趣"),
    Tool(func=tool2, name="tool2", description="曾经咨询过类似的产品,时，只调用这个工具，只返回return中的值"),
    Tool(func=search_db, name="search_db", description="询问2024年国家队时，搜索")
]





def get_tools_agent(user_input):
    template = '''Answer the following questions as best you can. You have access to the following tools:
    
               {tools}
    
               Use the following format:
    
               Question: the input question you must answer
               Thought: you should always think about what to do
               Action: the action to take, should be one of [{tool_names}]
               Action Input: the input to the action
               Observation: the result of the action
               ... (this Thought/Action/Action Input/Observation can repeat N times)
               Thought: I now know the final answer
               Final Answer: the final answer to the original input question
    
               Begin!
    
               Question: {input}
               Thought:{agent_scratchpad}'''

    prompt = ChatPromptTemplate.from_template(template)

    agent = create_react_agent(tools=tools, llm=llm, prompt=prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    ret = agent_executor.invoke({'input': user_input})
    return ret
